#!/bin/bash

# 脚本名称：expand_disk.sh

# 功能说明：
# 本脚本用于扩展虚拟磁盘或物理磁盘的最后一个分区的大小和文件系统大小。

# 主要流程包括：
# 1. 安装必要的依赖工具（如 qemu-img、qemu-nbd、parted 等）。
# 2. 根据用户输入判断设备是物理磁盘还是虚拟磁盘镜像。
# 3. 如果是虚拟磁盘镜像，则使用 qemu-nbd 挂载到一个空闲的 nbd 设备上。
# 4. 显示设备信息。
# 5. 扩展最后一个分区的大小，并扩展 ext2/3/4 文件系统以匹配新容量。
# 6. 使用 trap 在退出时自动断开 nbd 挂载（如果有的话）。

# 注意事项：
# - 仅支持 ext2/3/4 文件系统的分区扩展。
# - 需要 root 权限运行部分命令，如 qemu-nbd、parted、resize2fs。
# - 在执行分区扩展前，建议备份重要数据，以防万一。

# 安装依赖包
install_dependencies() {
    local required=("qemu-img" "qemu-nbd" "partprobe" "resize2fs" "parted")
    local missing=()

    for cmd in "${required[@]}"; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            missing+=("$cmd")
        fi
    done

    if [ ${#missing[@]} -eq 0 ]; then
        return 0
    fi

    echo "检测到缺少依赖命令: ${missing[*]}，尝试自动安装..."
    sudo apt-get update -qq
    echo "正在安装 qemu-utils parted e2fsprogs..."
    sudo apt-get install -y -qq qemu-utils parted e2fsprogs >/dev/null
}

# 全局变量
disk=""
nbd_device=""
target_device=""
virtual_disk_attached=false
mount_point="/mnt/temp_mount"

# 清理函数：如果虚拟磁盘已挂载，则断开 nbd 设备
cleanup() {
    if $virtual_disk_attached && [ -n "$nbd_device" ]; then
        echo "清理：断开 nbd 设备 $nbd_device"
        sudo qemu-nbd -d "$nbd_device"
    fi
    if mountpoint -q "$mount_point"; then
        sudo umount "$mount_point"
        sudo rmdir "$mount_point"
    fi
}
trap cleanup EXIT SIGINT SIGTERM

# 查找空闲的 nbd 设备
select_free_nbd() {
    for dev in /dev/nbd*; do
        if [ -z "$(lsblk -n "$dev" | awk 'NR>1')" ]; then
            echo "$dev"
            return 0
        fi
    done
    return 1
}

# 挂载虚拟磁盘
mount_virtual_disk() {
    sudo modprobe nbd max_part=8
    nbd_device=$(select_free_nbd)
    if [ -z "$nbd_device" ]; then
        echo "未找到可用的 nbd 设备，退出。"
        exit 1
    fi
    echo "将虚拟磁盘镜像挂载到 nbd 设备：$nbd_device"
    sudo qemu-nbd -c "$nbd_device" "$disk"
    sudo udevadm settle
    virtual_disk_attached=true
}

# 获取用户输入
get_disk_input() {
    if [ -n "$1" ]; then
        disk="$1"
    else
        echo "没有提供设备或镜像文件，现列出所有磁盘："
        lsblk -o NAME,MAJ:MIN,RM,LABEL,SIZE,FSTYPE,UUID,RO,TYPE,MOUNTPOINT,VENDOR,HOTPLUG -e7
        read -rp "请输入设备路径（如 /dev/sda）或虚拟磁盘镜像文件路径: " disk
    fi
}

# 处理设备类型
process_disk() {
    if [ -b "$disk" ]; then
        echo "选择的物理磁盘: $disk"
        target_device="$disk"
    elif [ -f "$disk" ]; then
        echo "选择的虚拟磁盘镜像文件: $disk"
        mount_virtual_disk
        target_device="$nbd_device"
    else
        echo "错误： '$disk' 既不是块设备也不是普通文件。"
        exit 1
    fi
}

# 显示设备信息
display_disk_info() {
    echo "显示设备信息："
    lsblk -o NAME,MAJ:MIN,RM,LABEL,SIZE,FSTYPE,UUID,RO,TYPE,MOUNTPOINT,VENDOR,HOTPLUG "$target_device"
}

# 获取分区的文件系统类型
get_fs_type() {
    local partition="$1"
    sudo blkid -o value -s TYPE "$partition"
}

# 获取指定设备的分区列表（仅显示分区，不显示磁盘设备）
get_partitions() {
    lsblk -nlo NAME,TYPE "$1" | grep -w "part" | awk '{print $1}'
}

# 获取指定设备的最后一个分区
get_last_partition() {
    local device="$1"
    local partitions
    partitions=$(get_partitions "$device")
    echo "$partitions" | tail -n1
}

# 扩展磁盘的最后一个分区
expand_last_partition() {
    local device="$1"
    local last_part
    local fs_type
    local part_num
    local disk_size_sectors
    local end_sector

    # 请求用户确认操作
    echo -n "确认要扩展最后一个分区到磁盘末尾吗？(Yes/No) "
    read -r answer
    if [ "$answer" != "Yes" ]; then
        echo "操作已取消"
        exit 0
    fi

    # 获取最后一个分区信息
    last_part=$(get_last_partition "$device")
    if [ -z "$last_part" ]; then
        echo "错误：未找到有效分区"
        exit 1
    fi
    last_part="/dev/${last_part}"

    # 获取并显示文件系统类型
    fs_type=$(get_fs_type "$last_part")
    echo "检测到最后一个分区 $last_part 的文件系统类型为: $fs_type"

    # 显示操作前的分区信息
    echo "操作前的分区信息："
    display_disk_info

    # 提取分区号
    part_num=$(echo "$last_part" | grep -o '[0-9]*$')
    if [ -z "$part_num" ]; then
        echo "错误：无法获取分区号"
        exit 1
    fi

    # 获取磁盘总大小（以扇区为单位）
    disk_size_sectors=$(sudo blockdev --getsz "$device")
    end_sector=$((disk_size_sectors - 1))  # 计算最后一个扇区号

    # 使用parted调整分区大小
    echo "调整分区 $part_num 到结束扇区 100%..."
    sudo parted -s "$device" unit s resizepart "$part_num" "100%"
    if [ $? -ne 0 ]; then
        echo "调整分区大小失败"
        exit 1
    fi

    # 更新内核分区表
    sudo partprobe "$device"
    if [ $? -ne 0 ]; then
        echo "刷新分区表失败"
        exit 1
    fi

    # 根据文件系统类型执行相应的扩展操作
    case $fs_type in
        ext4|ext3|ext2)
            echo "检查并扩展文件系统..."
            sudo e2fsck -f -y "$last_part"  # 检查文件系统
            if [ $? -ne 0 ]; then
                echo "文件系统检查失败"
                exit 1
            fi
            sudo resize2fs "$last_part"     # 扩展文件系统
            if [ $? -ne 0 ]; then
                echo "文件系统扩展失败"
                exit 1
            fi
            ;;
        *)
            echo "错误：不支持的文件系统类型 $fs_type"
            exit 1
            ;;
    esac

    # 显示操作后的分区信息
    echo -e "\n操作后的分区信息："
    display_disk_info
}

# 主函数
main() {
    install_dependencies
    get_disk_input "$1"
    process_disk
    display_disk_info
    expand_last_partition "$target_device"
}

main "$@"